import random
import numpy as np
import torch
from collections import deque

class ReplayBuffer:
    def __init__(self, capacity, num_agents):
        self.capacity = capacity
        self.num_agents = num_agents
        self.buffer = deque(maxlen=capacity)

    def add(self, obs, actions, rewards, next_obs, dones):
        self.buffer.append((obs, actions, rewards, next_obs, dones))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        obs_batch = []
        actions_batch = []
        rewards_batch = []
        next_obs_batch = []
        dones_batch = []
        for i in range(self.num_agents):
            obs_batch.append(torch.tensor(np.array([sample[0][i] for sample in batch]), dtype=torch.float32))
            actions_batch.append(torch.tensor(np.array([sample[1][i] for sample in batch]), dtype=torch.float32))
            rewards_batch.append(torch.tensor(np.array([sample[2][i] for sample in batch]), dtype=torch.float32).unsqueeze(1))
            next_obs_batch.append(torch.tensor(np.array([sample[3][i] for sample in batch]), dtype=torch.float32))
            dones_batch.append(torch.tensor(np.array([sample[4][i] for sample in batch]), dtype=torch.float32).unsqueeze(1))
        global_obs = torch.cat(obs_batch, dim=1)
        global_actions = torch.cat(actions_batch, dim=1)
        global_next_obs = torch.cat(next_obs_batch, dim=1)
        return obs_batch, actions_batch, rewards_batch, global_obs, global_actions, next_obs_batch, global_next_obs, dones_batch

    def __len__(self):
        return len(self.buffer)
